
from pylab import *
from sympy import Matrix, Function, var, expand, Symbol, symbols
from sympy.core.cache import clear_cache
from core import *
import re

# NOTAS:
# ya esta la correccion a la funcion adj, ahora es Adj. VSC Canela (20.05.2011)
#
# la funcion commutator la cambie a Commutator, para seguir el formato de 
#   funciones con mayuscula (4.09.2011)
#
# la funcion qmatrix cambio a Qmatrix para seguir el formato (4.09.2011)
#
# anexe las funciones TraceOut, Trace y Normalize; estan comentadas. (18.09.2011)
#
# anexe las funciones Dictionary, Qch, y Qeq; estan comentadas. (19.09.2011)
#
# funcion qeq es substituida por Qeq, pero se mantiene en este archivo
#   la qeq. (19.09.2011)
#
# la libreria pylab se debe importar antes que todas las demas, si no marca errores
#   (19.09.2011)
# 
# aniadida funcion SubsSol que sirve para graficar observables.
#   (1.03.2012)
# Agregue algo que limpia el cache, quitarlo cuando sympy lo haga automaticamente
# En subssol, cambio la sustitucion de obs(diccio) a obs.sub(diccio). Esto para que
# funcione con sage que tiene una version mas vieja de sympy

def X(A = 'default'):
    return Ket(0, A)*Bra(1, A) + Ket(1, A)*Bra(0, A)
def Y(A = 'default'):
    return -I*Ket(0, A)*Bra(1, A) + I*Ket(1, A)*Bra(0, A)
def Z(A = 'default'):
    return Ket(0, A)*Bra(0, A) - Ket(1, A)*Bra(1, A)

def Adj(u):
    """Returns the corresponding hermitian conjugate."""
    if isinstance(u, Bra):
        return Ket(u.eig, u.op)
    elif isinstance(u, Ket):
        return Bra(u.eig, u.op)
    elif isinstance(u, Operator):
        return AdjointOperator(u.op)
    elif isinstance(u, AdjointOperator):
        return op(u.op)
    elif isinstance(u, DiracSum):
        return DiracSum(*map(Adj, u.terms))
    elif isinstance(u, DiracMult):
        l = u.factors
	# la siguiente linea corrige el error que cambiaba la entrada original
	L = l[:]
        L.reverse()
        return DiracMult(u.coef.conjugate(), *map(Adj, L))
    else:
        return u.conjugate()

def Commutator(u, v):
    """Returns the commutator [u, v] of u and v."""
    return u*v - v*u

def Qmatrix(expression, basis = 'default'):
    """This function returns the matrix representation of any
    operator, e.g. |0><1| + |1><0|. The second argument is the
    basis of the matrix, given as a list of vectors (kets), e.g.
    [Ket(0), Ket(1)]. If no basis is specified, it is first found out
    how many different kets or bras the expression contains, to determine
    the basis and the size of the matrix. For the moment, these can only
    be eigenvectors of a single operator."""
    expression = DiracSum(expression)
    if basis == 'default':
        eigs = []; ops = []
        for term in expression.terms:
            if isinstance(term, DiracMult):
                pass
            else:
                term = DiracMult(1, term)
            for fact in term.factors:
                if isinstance(fact, Ket) or isinstance(fact, Bra):
                    eigs.append(fact.eig)
                    ops.append(fact.op)
        eigs = list(set(eigs)); eigs.sort()
        ops = list(set(ops)); ops.sort()
        if len(ops) > 1:
            try:
                raise Exception, 'Matrix form can only be obtained for operators \
                                  acting on a single Hilbert space.'
            except Exception:
                print 'An exception occurred.'
                raise
        else:
            oper = ops[0]
        M = []
        for i in eigs:
            row = []
            for j in eigs:
                row.append(Bra(i,oper)*expression*Ket(j,oper))
            M.append(row)
        M = Matrix(M)
    else:
        M = []
        for i in basis:
            row = []
            for j in basis:
                row.append(Adj(i)*expression*j)
            M.append(row)
        M = Matrix(M)
    return M

def qeq(expression, *variables):
    """This function converts a quantum expression into a symbolic
    expression when this operation has sense. For instance,
    <k,l|A|m,n> transforms into A_k_l_m_n, a function of the
    variables specified in *variables.
    Currently works with only one variable"""
    if len(variables) > 1:
        raise Exception, 'Sorry, qeq currently works with at most \
        one variable.'
    elif len(variables) == 1:
        variable = var(str(variables[0]))
    expression = DiracSum(expression)
    for term in expression.terms:
        if isinstance(term, DiracMult):
            pass
        else:
            term = evaluar(DiracMult(1, term))
    expression2 = 0
    for term in expression.terms:
        if isinstance(term, DiracMult):
            term2 = term.coef
            for n in range(len(term.factors)):
                if isinstance(term.factors[n], Operator):
                    # s is the name of the function or symbolic var,
                    # supposing that the operator's label is a string:
                    if isinstance(term.factors[n].op, str):
                        s = term.factors[n].op
                    else:
                        raise Exception, "In function qeq, the operator's \
                        label must be a string."
                    kets = []; bras = []
                    layer = 1 # i.e. the current bra & ket layer
                    another_ket = True
                    while another_ket:
                        if (isinstance(term.factors[n-layer], Bra) and
                        isinstance(term.factors[n+layer], Ket)):
                            if (term.factors[n-layer].op ==
                            term.factors[n+layer].op):
                                bras.append(term.factors[n-layer].op)
                                kets.append(term.factors[n+layer].op)
                        layer += 1
                        if n+layer < len(term.factors):
                            if not(isinstance(term.factors[n+layer], Ket)):
                                another_ket = False
                        else:
                            another_ket = False
                    for i in range(len(bras)):
                        s += "_%s" % (str(bras[i]))
                    for i in range(len(kets)):
                        s += "_%s" % (str(kets[i]))
                    if len(variables) > 0:
                        term2 = term2 * Function(s)(variable)
                    else:
                        term2 = term2 * var(s)
                elif isinstance(term.factors[n], AdjointOperator):
                    if isinstance(term.factors[n].op, str):
                        s = term.factors[n].op
                    else:
                        raise Exception, "In function qeq, the operator's \
                        label must be a string."
                    kets = []; bras = []
                    layer = 1 # i.e. the current bra & ket layer
                    another_ket = True
                    while another_ket:
                        if (isinstance(term.factors[n-layer], Bra) and
                        isinstance(term.factors[n+layer], Ket)):
                            if (term.factors[n-layer].op ==
                            term.factors[n+layer].op):
                                bras.append(term.factors[n+layer].op)
                                kets.append(term.factors[n-layer].op)
                        layer += 1
                        if n+layer < len(term.factors):
                            if not(isinstance(term.factors[n+layer], Ket)):
                                another_ket = False
                        else:
                            another_ket = False
                    for i in range(len(bras)):
                        s += "_%s" % (str(bras[i]))
                    for i in range(len(kets)):
                        s += "_%s" % (str(kets[i]))
                    if len(variables) > 0:
                        term2 = term2 * Function(s)(variable)
                    else:
                        term2 = term2 * var(s)
        else:
            term2 = term
        expression2 = expression2 + term2
    if isinstance(expression2, DiracSum):
        try:
            raise Exception, 'Could not create symbolic expression. The \
            only allowed quantum objects must be of the form \
            <b1, b2, ..., bn|A|k1, k2, ..., kn>, \
            where A is an operator or a dual operator.'
        except Exception:
            print 'An exception occurred.'
            raise
    return expression2

def TraceOut(psi, a):
    """This function computes the partial trace of a series of exterior products,
    for example: |a_1>|b_2><c_1|<d_2| where 1 and 2 mean different Hilbert spaces.
    psi is the complete expression, and a is the tag of the Hilbert space to be
    traced out."""
    psi = expandir(psi)
    Lterms = psi.terms
    htags = []
    elements = []
    l = []
    for i in Lterms:
        Lfactors = i.factors
        for j in Lfactors:
            htags.append(j.op)
            elements.append(j)
    htags = list(set(htags))
    d = {}
    for i in htags:
        d[i] = []
    for i in d:
        for j in elements:
            if j.op == i:
                d[i].append(j.eig)
    for i in d:
        d[i] = list(set(d[i]))
    term = 0
    for i in d.keys():
        if i == a:
            for j in d[i]:
                term = term + Bra(j,i) * psi * Ket(j,i)
            psi = term
            term = 0
    return psi

def Trace(psi, basis = 'default'):
    """Computes the total trace of a sum of exterior products. First it finds the
    total number of Hilbert spaces and their tags, and then it uses TraceOut over
    all of them succesively, returning a complex."""
    if basis == 'default':
        psi = expandir(psi)
        Lterms = psi.terms
        htags = []
        for i in Lterms:
            Lfactors = i.factors
            for j in Lfactors:
                htags.append(j.op)
        htags = list(set(htags))
        for i in htags:
            psi = TraceOut(psi, i)
    else:
        temp = 0
        for l in basis:
            temp = temp + Adj(l)*psi*l
        psi = temp
    return psi

def Normalize(psi):
    """Produces a vector in a Hilbert space with norm equal to unity"""
    norm = (Adj(psi)*psi)
    norm = float(norm)
    psi = psi / (norm**0.5)
    return psi

def Dictionary(x, R):
    """This function creates a dictionary object with 'words' as strings of
    matrix elements of the operator R, in the base x. x is given as a list
    containing all the elements of the basis, and R is necessarily defined
    as an Operator. The function creates variables representing a complex
    number, e.g. w + I*y, where w and y are defined as real variables. If the
    base used has n elements, n*(n+1)/2 variables are created, asigning each
    'word' a variable. Complex conjugation is taken into account as follows:
    if <1|R|2> corresponds with w + I*y then <2|R|1> corresponds with w - I*y.
    Also diagonal elements are real, i.e. <a|R|a> corresponds with some w real.
    The 'definitions' of each 'word' are the variables themselves."""
    n = len(x)
    u = []
    w = []
    t = []
    for i in range( (n*(n+1))/2 ):
        u.append('Re%d' %i)
        u[i] = symbols('Re%d' %i, real=True, each_char=False)
        w.append('Im%d' %i)
        w[i] = symbols('Im%d' %i, real=True, each_char=False)
    for i in range(len(u)):
        t.append(u[i] + I*w[i])
    D = {}
    k = 0
    for i in range(n):
        for j in range(n):
            if i <= j:
                D[str(Adj(x[i])*R*x[j])] = t[k]
            else:
                D[str(Adj(x[i])*R*x[j])] = conjugate(D[str(Adj(x[j])*R*x[i])])
            if i <= j:
                k = k + 1
    return D


def DictionaryN(x, R):
    """This function creates a dictionary object with 'words' as strings of
    matrix elements of the operator R, in the base x. x is given as a list
    containing all the elements of the basis, and R is necessarily defined
    as an Operator. The function creates variables representing a complex
    number, e.g. w + I*y, where w and y are defined as real variables. If the
    base used has n elements, n*n variables are created, asigning each
    'word' a variable. Is is assumed that all variables al complex.
    The 'definitions' of each 'word' are the variables themselves."""
    n = len(x)
    u = []
    w = []
    t = []
    for i in range( n*n ):
        u.append('Re%d' %i)
        u[i] = symbols('Re%d' %i, real=True, each_char=False)
        w.append('Im%d' %i)
        w[i] = symbols('Im%d' %i, real=True, each_char=False)
    for i in range(len(u)):
        t.append(u[i] + I*w[i])
    D = {}
    k = 0
    for i in range(n):
        for j in range(n):
            D[str(Adj(x[i])*R*x[j])] = t[k]
            k = k + 1
    return D


def Qch(x, D):
    """This function looks within the expression x, for 'words' in the 
    dictionary D, and substitutes them with their 'definition'. It 
    discriminates coefficients and other terms which are not in D, leaving
    them the same."""
    if x == 0.0:
        l = [0]
    else:
        x = DiracSum(x)
        L = x.terms
        l = []
        for i in L:
            y = DiracMult(1,*i.factors)
            if str(y) in D:
                y = str(y)
                y = D[y]
            l.append(y * (i.coef))
    return Add(*l)

def QeqN(R, Rdot, b, name, nameD):
    """This function creates a file called name, which defines a function
    in the sintaxis needed to feed scipy's ODE solver, odeint. Its main 
    purpose is to study time evolution of an Operator R, with an expression
    for its time derivative Rdot. b denotes the base, as a list of elements
    of a base.
    It first creates a dictionary D with the function Dictionary, using the list b,
    and the Operator R, then it finds the different expressions of time evolution 
    for every matrix element substituting them for variables using Qch. Finally it 
    writes the function f to name, where the ODE has the form dy/dt = f(y,t).
    y should be understood as a vector, and t the variable 'time'. Thus y is defined
    entry-like as a list e.g. y[0]=x, y[1]=w, etc.
    f returns also a list, containig the expressions of the derivatives of each entry
    of y as entries of another list. 
    odeint understands this as a collection of coupled ODE to be solved, specifying
    an initial condition and the time step."""
    D = Dictionary(b, R)
    #print D
    f = open(name, 'w')
    f.write('from __future__ import division\n')
    f.write('from sympy import var, I\n')
    f.write('def f(y, t):\n')
    
    # crea la lista con las variables Re, Im
    M, u, w = [], [], []
    n = len(b)
    for i in range( (n*(n+1))/2 ):
        u.append('Re%d' %i)
        u[i] = symbols('Re%d' %i, real=True, each_char=False)
        w.append('Im%d' %i)
        w[i] = symbols('Im%d' %i, real=True, each_char=False)
    k = 0
    for i in range(n):
        for j in range(n):
            if i <= j:
                M.append(u[k])
                M.append(w[k])
                k = k + 1
    # termina de crearla

    # crea la lista con los strings y[0], y[1]
    L = ['y[' + '%d' %i + ']' for i in range(n*(n+1))]
    # termina de crearla

    # Crea el diccionario para hacer la substitucion de forma rapida
    lista={}
    for mi in range((n*(n+1))):
        lista[M[mi]]=Symbol(L[mi])

    
    f.write('    ' + 'return [')    
    for numero,i in enumerate(b):
        print numero,"/",len(b)
        clear_cache()   #Limpia el cache para que no ocupe toda la memoria
        for j in b:
            if b.index(i) <= b.index(j):
                z = (Adj(i)*Rdot*j + Adj(j)*Rdot*i)/2.
                zz = Qch(z, D)
                zz = expand(zz,complex=True)
                for var_subs in zz.atoms(Symbol):
                    if lista.has_key(var_subs):
                        zz = zz.subs(var_subs,lista[var_subs])
                f.write( str(zz)+',\n')         
                z = (Adj(i)*Rdot*j - Adj(j)*Rdot*i)/(2.*I)
                zz = Qch(z, D)
                zz = expand(zz,complex=True)
                for var_subs in zz.atoms(Symbol):
                    if lista.has_key(var_subs):
                        zz = zz.subs(var_subs,lista[var_subs])
                f.write(str(zz))     
                if b.index(i)*b.index(j)==(len(b)-1)**2:
                    f.write(']\n')
                else:
                    f.write(',\n')
    # el diccionario que relaciona '<i|R|j>' con las y[i]
    # realmente con las "var('y[i]')"

####    EMPIEZA VIEJO
    #for conta1 in D.keys():
    #    for conta2 in range(n**2):
    #        D[conta1] = D[conta1].subs(M[conta2], Symbol(L[conta2])) #cambie a LL
    #for conta3 in D.keys():
    #    D[conta3] = str(D[conta3])
    # termina de crearlo
    #print D
####    TERMINA VIEJO    
 
#    LL = ["var('y" + '%d' %i + "',real=True)" for i in range(n**2)]
    f.write(nameD + '=' + '{')
    for count1 in b:
        for count2 in b:
                temp = Adj(count1)*R*count2
                temp = Qch(temp, D)
                for var_subs in temp.atoms(Symbol):
                    etiqueta=re.search("\[(.*)\]",str(lista[var_subs]))
                    temp = temp.subs(var_subs,Symbol("var('y"+etiqueta.groups()[0]+"',real=True)"))
                f.write("'" + str(Adj(count1)*R*count2) + "'" + ':' + \
                str(temp) + ',')
    f.write( '}')
    f.close()

def QeqNN(R, Rdot, b, name, nameD):
    """This function creates a file called name, which defines a function
    in the sintaxis needed to feed scipy's ODE solver, odeint. Its main 
    purpose is to study time evolution of an Operator R, with an expression
    for its time derivative Rdot. b denotes the base, as a list of elements
    of a base.
    It first creates a dictionary D with the function Dictionary, using the list b,
    and the Operator R, then it finds the different expressions of time evolution 
    for every matrix element substituting them for variables using Qch. Finally it 
    writes the function f to name, where the ODE has the form dy/dt = f(y,t).
    y should be understood as a vector, and t the variable 'time'. Thus y is defined
    entry-like as a list e.g. y[0]=x, y[1]=w, etc.
    f returns also a list, containig the expressions of the derivatives of each entry
    of y as entries of another list. 
    odeint understands this as a collection of coupled ODE to be solved, specifying
    an initial condition and the time step."""
    D = DictionaryN(b, R)
    #print D
    f = open(name, 'w')
    f.write('from __future__ import division\n')
    f.write('from sympy import var, I\n')
    f.write('def f(y, t):\n')
    
    # crea la lista con las variables Re, Im
    M, u, w = [], [], []
    n = len(b)
    for i in range( n*n ):
        u.append('Re%d' %i)
        u[i] = symbols('Re%d' %i, real=True, each_char=False)
        w.append('Im%d' %i)
        w[i] = symbols('Im%d' %i, real=True, each_char=False)
    k = 0
    for i in range(n):
        for j in range(n):
            M.append(u[k])
            M.append(w[k])
            k = k + 1
    # termina de crearla

    # crea la lista con los strings y[0], y[1]
    L = ['y[' + '%d' %i + ']' for i in range(2*n*n)]
    # termina de crearla

    # Crea el diccionario para hacer la substitucion de forma rapida
    lista={}
    for mi in range(2*n*n):
        lista[M[mi]]=Symbol(L[mi])

    f.write('    ' + 'return [')    
    for numero,i in enumerate(b):
        print numero,"/",len(b)
        clear_cache()   #Limpia el cache para que no ocupe toda la memoria
        for j in b:
            z = Adj(i)*Rdot*j
            zz = Qch(z, D)
            zz = expand((zz+conjugate(zz))/2.0,complex=True)
            for var_subs in zz.atoms(Symbol):
                if lista.has_key(var_subs):
                    zz = zz.subs(var_subs,lista[var_subs])
            f.write( str(zz)+',\n')         
            zz = Qch(z, D)
            zz = expand((zz-conjugate(zz))/(2.0*I),complex=True)
            for var_subs in zz.atoms(Symbol):
                if lista.has_key(var_subs):
                    zz = zz.subs(var_subs,lista[var_subs])
            f.write(str(zz))     
            if b.index(i)*b.index(j)==(len(b)-1)**2:
                f.write(']\n')
            else:
                f.write(',\n')
    # el diccionario que relaciona '<i|R|j>' con las y[i]
    # realmente con las "var('y[i]')"

####    EMPIEZA VIEJO
    #for conta1 in D.keys():
    #    for conta2 in range(n**2):
    #        D[conta1] = D[conta1].subs(M[conta2], Symbol(L[conta2])) #cambie a LL
    #for conta3 in D.keys():
    #    D[conta3] = str(D[conta3])
    # termina de crearlo
    #print D
####    TERMINA VIEJO    
 
#    LL = ["var('y" + '%d' %i + "',real=True)" for i in range(n**2)]
    f.write(nameD + '=' + '{')
    for count1 in b:
        for count2 in b:
            temp = Adj(count1)*R*count2
            temp = Qch(temp, D)
            for var_subs in temp.atoms(Symbol):
                etiqueta=re.search("\[(.*)\]",str(lista[var_subs]))
                temp = temp.subs(var_subs,Symbol("var('y"+etiqueta.groups()[0]+"',real=True)"))
            f.write("'" + str(Adj(count1)*R*count2) + "'" + ':' + \
            str(temp) + ',')
    f.write( '}')
    f.close()




def SubsSol(obs, sol):
    """This function converts variable expressions to the time evolution
    (obtained by odeint) of the corresponding observable. The input expression
    is obs, while sol is the object returned by odeint, and b is the base
    used in QeqN"""
    temp={}; L = []
    if len(obs.args) == 0:
        for count1 in range(len(sol)):
            L.append(sol[:, int(str(obs).partition('y')[-1])][count1])
            clear_cache()   #Limpia el cache para que no ocupe toda la memoria
    else:
        for count1 in range(len(sol)):
            List = obs.args
            for count2 in List:
                for count3 in count2.atoms(Symbol):
                    temp[count3] = sol[:, int(str(count3).partition('y')[-1])][count1]
                    clear_cache()   #Limpia el cache para que no ocupe toda la memoria
            L.append(obs.subs(temp))
            clear_cache()   #Limpia el cache para que no ocupe toda la memoria
    L=array(L)
    return L


def InitialCondition(R,obs, base, dic):
    """This function produces an initial condition in the form
    of a list (ordered in accordance to the variables yi), given
    the initial condition as a sum of exterior products (obs). The
    base (base) has to be given as a list of kets, as well as the dictionary (dic)
    generated by QeqN. The density operator (R) used in QeqN has to be given
    as well."""
    n=len(base)
    ini = [0]*(n*(n+1))
    tempdic = {}
    for numero,i in enumerate(base):
        print numero,"/",len(base)
        for j in base:
            if base.index(i) <= base.index(j):
                tempvar = Qch(Adj(i)*R*j,dic)
                tempdic[expand( (tempvar + conjugate(tempvar))/2. )] = expand((Adj(i)*obs*j + Adj(j)*obs*i)/2.)
                tempdic[expand( (tempvar - conjugate(tempvar))/(2.*I) )] = expand((Adj(i)*obs*j - Adj(j)*obs*i)/(2.*I))
    for i in tempdic.keys():
        ini[int(str(i).partition('y')[-1])] = tempdic[i]
    return ini

def InitialConditionN(R,obs, base, dic):
    """This function produces an initial condition in the form
    of a list (ordered in accordance to the variables yi), given
    the initial condition as a sum of exterior products (obs). The
    base (base) has to be given as a list of kets, as well as the dictionary (dic)
    generated by QeqN. The density operator (R) used in QeqN has to be given
    as well."""
    n=len(base)
    ini = [0]*(2*n*n)
    tempdic = {}
    for numero,i in enumerate(base):
        print numero,"/",len(base)
        for j in base:
            tempvar = Qch(Adj(i)*R*j,dic)
            tempdic[expand( (tempvar + conjugate(tempvar))/2. )] = expand((Adj(i)*obs*j + Adj(j)*obs*i)/2.)
            tempdic[expand( (tempvar - conjugate(tempvar))/(2.*I) )] = expand((Adj(i)*obs*j - Adj(j)*obs*i)/(2.*I))
    for i in tempdic.keys():
        ini[int(str(i).partition('y')[-1])] = tempdic[i]
    return ini
        









